import os
import random
from functools import partial
from multiprocessing.pool import Pool
from pathlib import Path

import cv2
import pandas as pd
from tqdm import tqdm

from constants import DFDC
from option import parse_args
from preprocessing.generate_folds import cmp_frame_name, get_real_fake_pairs, cmp, real_fake_split

os.environ["MKL_NUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
os.environ["OMP_NUM_THREADS"] = "1"

cv2.ocl.setUseOpenCL(False)
cv2.setNumThreads(0)

SAMPLE_INTERVALS = [-40, -30, -20, -10, 0, 10, 20, 30, 40]


def check_crops2(ori_dir, fake_dir, frame):
    frame_split = frame.split('_')
    frame_num = int(frame_split[0])
    for interval in SAMPLE_INTERVALS:
        frame_cur = frame_num + interval
        img_name = f'{frame_cur}_{frame_split[1]}'
        ori_img_path = os.path.join(ori_dir, img_name)
        fake_img_path = os.path.join(fake_dir, img_name)
        if not os.path.exists(ori_img_path) or not os.path.exists(fake_img_path):
            return False
    return True


def get_paths(vid, label, root_dir, sample_num: int):
    part_id, ori_id, fake_id = vid
    ori_dir = os.path.join(root_dir, 'crops', part_id, ori_id)
    fake_dir = os.path.join(root_dir, 'crops', part_id, fake_id)
    mask_dir = os.path.join(root_dir, 'diffs', part_id, fake_id)

    data = []
    # Some unaligned videos have been removed
    if label == 1 and not os.path.exists(fake_dir):
        return data
    # Some masks may have been removed
    if label == 1 and not os.path.exists(mask_dir):
        return data

    frames = os.listdir(fake_dir)
    frames.sort(key=cmp_frame_name)
    for frame in frames:
        frame_split = str(frame).split('_')
        frame_num = int(frame_split[0])
        if frame_num % 50 != 0 or not check_crops2(ori_dir, fake_dir, frame):
            continue
        ori_img_path = os.path.join(ori_dir, frame)
        fake_img_path = os.path.join(fake_dir, frame)
        if label == 0 and os.path.exists(ori_img_path):
            data.append([ori_img_path, label, ori_id])
        elif label == 1 and os.path.exists(fake_img_path) and os.path.exists(ori_img_path):
            mask_path = os.path.join(mask_dir, "{}_diff.png".format(frame[:-4]))
            if os.path.exists(mask_path):
                data.append([fake_img_path, label, ori_id])

    return data


def save_folds(args, save_path, pairs_mode, mode, sample_num=100):
    data = []
    if mode == 'train_real':
        ori_ori_pairs = set([(part, ori, ori) for part, ori, fake in pairs_mode])
        ori_fake_pairs = set()
    elif mode == 'train_fake':
        ori_ori_pairs = set()
        ori_fake_pairs = set(pairs_mode)
    else:
        ori_ori_pairs = set([(part, ori, ori) for part, ori, fake in pairs_mode])
        ori_fake_pairs = set(pairs_mode)

    ori_ori_pairs = list(ori_ori_pairs)
    ori_fake_pairs = list(ori_fake_pairs)
    ori_ori_pairs.sort(key=cmp)
    ori_fake_pairs.sort(key=cmp)

    with Pool(processes=1) as p:
        # original label=0
        with tqdm(total=len(ori_ori_pairs)) as pbar:
            func = partial(get_paths, label=0, root_dir=args.root_dir, sample_num=sample_num)
            for v in p.imap_unordered(func, ori_ori_pairs):
                pbar.update()
                data.extend(v)
        # fake label=1
        with tqdm(total=len(ori_fake_pairs)) as pbar:
            func = partial(get_paths, label=1, root_dir=args.root_dir, sample_num=sample_num)
            for v in p.imap_unordered(func, ori_fake_pairs):
                pbar.update()
                data.extend(v)
    fold_data = []
    for img_path, label, ori_vid in data:
        path = Path(img_path)
        part = path.parent.parent.name
        video = path.parent.name
        file = path.name
        fold_data.append([os.path.join(part, video), file, label, os.path.join(part, ori_vid), int(file.split("_")[0])])
    # random.shuffle(fold_data)
    columns = ["video", "file", "label", "original", "frame"]
    file_name = '{}_st_{}.csv'.format(DFDC, mode)
    save_file = os.path.join(save_path, file_name)
    pd.DataFrame(fold_data, columns=columns).to_csv(save_file, index=False)


def main():
    random.seed(110)
    args = parse_args()
    save_path = '../data/{}'.format(DFDC)
    os.makedirs(save_path, exist_ok=True)

    # original_path = os.path.join(args.root_dir, 'original_sequences', 'youtube', 'c23')
    # manipulated_path = os.path.join(args.root_dir, 'manipulated_sequences', args.fake_type, 'c23')
    pairs_train = get_real_fake_pairs(args.root_dir, 'train', sample_num=3000)
    pairs_val = get_real_fake_pairs(args.root_dir, 'val', sample_num=1000)
    pairs_test = get_real_fake_pairs(args.root_dir, 'test', sample_num=1000)
    print(len(pairs_train), len(pairs_val), len(pairs_test))

    sample_num_train = 10
    sample_num_val = 10
    sample_num_test = 20
    check_mask = True

    save_folds(args, save_path, pairs_train, 'train', sample_num=sample_num_train)
    save_folds(args, save_path, pairs_val, 'val', sample_num=sample_num_val)
    save_folds(args, save_path, pairs_test, 'test', sample_num=sample_num_test)

    ori_fake_pairs_real, ori_fake_pairs_fake = real_fake_split(pairs_train)
    save_folds(args, save_path, ori_fake_pairs_real, 'train_real', sample_num=sample_num_train)
    save_folds(args, save_path, ori_fake_pairs_fake, 'train_fake', sample_num=sample_num_train)


if __name__ == '__main__':
    # Total: 2000
    main()
